# encoding=utf-8
from pydot import Node, Edge, Dot
class Dotter(Dot, Node, Edge):
	def __init__(self, tree, solution, initial, country):
# 		self.tree_dot=[k for k,v in itertools.groupby(sorted(tree))]
		self.tree_dot = tree
		self.solution = solution
		self.initial = initial
		self.graph = Dot(graph_type='digraph')
		self.country = country

	def write_dot_file(self):
		added = []
		sol = []
		for i in self.solution:
			sol.append(i)
		sol.insert(0, self.initial)
		number = 1
		for node in self.tree_dot:
			noh = node[0]
			if noh not in added:
				if noh in sol:
					number += 1
					self.graph.add_node(Node(noh, style='filled', fillcolor='green',
											label=noh.strip('0123456789').replace('_', ' ').title() + ' ' + str(number)))
				else:
					self.graph.add_node(Node(noh,
											label=noh.strip('0123456789').replace('_', ' ').title() + ' ' + str(number)))
			added.append(noh)


			for subnode in node[1]:
				if subnode not in added:
					number += 1
					if subnode in sol:
						self.graph.add_node(Node(subnode, style='filled', fillcolor='green',
												label=subnode.strip('0123456789').replace('_', ' ').title() + ' ' + str(number)))
					else:
						self.graph.add_node(Node(subnode,
												label=subnode.strip('0123456789').replace('_', ' ').title() + ' ' + str(number)))
				added.append(subnode)

		for node in self.tree_dot:
			noh = node[0]
			for subnode in node[1]:
				try:
					self.graph.add_edge(Edge(self.graph.get_node(noh), self.graph.get_node(subnode),
											label=str(self.country.get(noh.strip('0123456789'), subnode.strip('0123456789')))))
				except:
					print 'Erro no nó:', noh, subnode

		self.graph.write('graph.dot')

